import torch
import tkinter as tk
from tkinter import messagebox
import numpy as np

class GobangGame:
    def __init__(self, root):
        self.root = root
        self.root.title("PyTorch 五子棋")
        
        # 棋盘参数
        self.size = 15  # 15x15棋盘
        self.cell_size = 40  # 每个格子大小
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        # 初始化棋盘张量（0:空, 1:黑棋, 2:白棋）- PyTorch张量替代列表
        self.board = torch.zeros((self.size, self.size), dtype=torch.int8, device=self.device)
        self.current_player = 1  # 1:黑方(玩家), 2:白方(AI)
        self.game_over = False
        self.history = []  # 悔棋历史（存储张量状态）
        
        # 创建画布
        self.canvas = tk.Canvas(
            root, 
            width=self.size * self.cell_size, 
            height=self.size * self.cell_size,
            bg="#EEDC82"  # 棋盘底色（浅黄色）
        )
        self.canvas.pack(padx=10, pady=10)
        self.draw_board()
        
        # 绑定鼠标点击事件
        self.canvas.bind("<Button-1>", self.on_click)
        
        # 按钮面板
        self.panel = tk.Frame(root)
        self.panel.pack(pady=5)
        
        self.regret_btn = tk.Button(self.panel, text="悔棋", command=self.regret)
        self.regret_btn.pack(side=tk.LEFT, padx=5)
        
        self.restart_btn = tk.Button(self.panel, text="重新开始", command=self.restart)
        self.restart_btn.pack(side=tk.LEFT, padx=5)

        # 初始化AI评分网络（简单MLP）
        self.ai_net = self.init_ai_net()

    # 初始化AI评分网络（输入：棋盘状态特征，输出：落子得分）
    def init_ai_net(self):
        net = torch.nn.Sequential(
            torch.nn.Linear(self.size * self.size * 2, 256),  # 输入：黑白棋各自的占位情况（15*15*2）
            torch.nn.ReLU(),
            torch.nn.Linear(256, 128),
            torch.nn.ReLU(),
            torch.nn.Linear(128, 1)  # 输出：落子得分
        ).to(self.device)
        return net

    # 绘制棋盘
    def draw_board(self):
        # 画横线竖线
        for i in range(1, self.size):
            self.canvas.create_line(
                i * self.cell_size, 0, 
                i * self.cell_size, self.size * self.cell_size,
                fill="#000000"
            )
            self.canvas.create_line(
                0, i * self.cell_size, 
                self.size * self.cell_size, i * self.cell_size,
                fill="#000000"
            )
        # 画天元和星位
        star_positions = [(3,3), (3,11), (7,7), (11,3), (11,11)]
        for (x, y) in star_positions:
            self.canvas.create_oval(
                x*self.cell_size-3, y*self.cell_size-3,
                x*self.cell_size+3, y*self.cell_size+3,
                fill="#000000"
            )

    # 落子（PyTorch张量操作）
    def place_piece(self, x, y, player):
        # 转换为棋盘坐标
        col = round(x / self.cell_size)
        row = round(y / self.cell_size)
        
        # 边界判断 + 空位判断（张量索引操作）
        if 0 <= col < self.size and 0 <= row < self.size and self.board[row, col] == 0:
            # 保存当前棋盘状态（用于悔棋）
            self.history.append(self.board.clone())
            
            # 落子（张量赋值）
            self.board[row, col] = player
            
            # 绘制棋子
            center_x = col * self.cell_size
            center_y = row * self.cell_size
            color = "black" if player == 1 else "white"
            self.canvas.create_oval(
                center_x - self.cell_size//2 + 2,
                center_y - self.cell_size//2 + 2,
                center_x + self.cell_size//2 - 2,
                center_y + self.cell_size//2 - 2,
                fill=color, outline="black"
            )
            
            # 判定胜负（张量加速判断）
            if self.check_win(row, col, player):
                winner = "黑方" if player == 1 else "白方"
                messagebox.showinfo("游戏结束", f"{winner}获胜！")
                self.game_over = True
            elif torch.all(self.board != 0):  # 棋盘满了（张量全非零判断）
                messagebox.showinfo("游戏结束", "平局！")
                self.game_over = True
            else:
                # 切换玩家（AI落子）
                if player == 1:
                    self.ai_move()
            return True
        return False

    # 胜负判定（基于张量的4方向连子检测）
    def check_win(self, row, col, player):
        # 转换为numpy数组方便遍历（小规模数据CPU操作更快）
        board_np = self.board.cpu().numpy()
        directions = [(0,1), (1,0), (1,1), (1,-1)]  # 横、竖、斜4方向
        
        for dx, dy in directions:
            count = 1
            # 正向计数
            r, c = row + dx, col + dy
            while 0 <= r < self.size and 0 <= c < self.size and board_np[r, c] == player:
                count += 1
                r += dx
                c += dy
            # 反向计数
            r, c = row - dx, col - dy
            while 0 <= r < self.size and 0 <= c < self.size and board_np[r, c] == player:
                count += 1
                r -= dx
                c -= dy
            if count >= 5:
                return True
        return False

    best_score = -float("inf")
         best_pos = None
         
         # 生成棋盘特征（黑白棋各自的二进制特征）
         black_feature = (self.board == 1).float().flatten()
         white_feature = (self.board == 2).float().flatten()
         board_feature = torch.cat([black_feature, white_feature], dim=0).unsqueeze(0)
         
         # 遍历所有空位，用网络评分
         for row in range(self.size):
             for col in range(self.size):
                 if self.board[row, col] == 0:
                     # 模拟落子
                     self.board[row, col] = 2
                     
                     # 生成落子后的特征
                     black_feature_temp = (self.board == 1).float().flatten()
                     white_feature_temp = (self.board == 2).float().flatten()
                     feature_temp = torch.cat([black_feature_temp, white_feature_temp], dim=0).unsqueeze(0)
                     
                     # 网络评分（禁用梯度计算，加速推理）
                     with torch.no_grad():
                         score = self.ai_net(feature_temp).item()
                     
                     # 恢复棋盘
                     self.board[row, col] = 0
                     
                     # 记录最优位置
                     if score > best_score:
                         best_score = score
                         best_pos = (row, col)
         
         # 下最优位置
         if best_pos:
             row, col = best_pos
             self.place_piece(col * self.cell_size, row * self.cell_size, 2)
     # 悔棋（张量克隆恢复）
     def regret(self):
         if self.history and not self.game_over:
             # 恢复上一步棋盘状态
             self.board = self.history.pop()
             # 重新绘制棋盘
             self.canvas.delete("all")
             self.draw_board()
             # 重新绘制所有棋子（基于张量状态）
             for row in range(self.size):
                 for col in range(self.size):
                     player = self.board[row, col].item()
                     if player != 0:
                         color = "black" if player == 1 else "white"
                         center_x = col * self.cell_size
                         center_y = row * self.cell_size
                         self.canvas.create_oval(
                             center_x - self.cell_size//2 + 2,
                             center_y - self.cell_size//2 + 2,
                             center_x + self.cell_size//2 - 2,
                             center_y + self.cell_size//2 - 2,
                             fill=color, outline="black"
                         )
             # 恢复玩家回合
             self.current_player = 1
     # 重新开始（张量重置）
     def restart(self):
         self.board = torch.zeros((self.size, self.size), dtype=torch.int8, device=self.device)
         self.current_player = 1
         self.game_over = False
         self.history = []
         self.canvas.delete("all")
         self.draw_board()
     # 鼠标点击事件
     def on_click(self, event):
         if not self.game_over and self.current_player == 1:
             self.place_piece(event.x, event.y, 1)
 if __name__ == "__main__":
     root = tk.Tk()
     game = GobangGame(root)
     root.mainloop()